package com.nl.kafka;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Properties;
import java.util.concurrent.atomic.AtomicBoolean;

import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.nl.util.StringUtil;



/**
 * KafkaConsumer
 * @author zhxf
 *
 */
public class KafkaConsumer {//com.kafka.KafkaConsumer

	private static final Logger LOG = LoggerFactory.getLogger(KafkaConsumer.class);

	String brokerList;
	String groupId;
	org.apache.kafka.clients.consumer.KafkaConsumer<String, byte[]> consumer;
	List<String> topic;
	AtomicBoolean rebalanceFlag;
	long batchDurationMillis;

	public KafkaConsumer(String brokerList, String groupId) {
		this.brokerList = brokerList;
		this.groupId = groupId;
		rebalanceFlag = new AtomicBoolean(false);
	}

	public void setTopic(List<String> topic) {
		this.topic = topic;
	}
	
	public void setBatchDurationMillis(long batchDurationMillis){
		this.batchDurationMillis = batchDurationMillis;
	}

	public void connect() {
		consumer = new org.apache.kafka.clients.consumer.KafkaConsumer<String, byte[]>(getConsumerProps());
		consumer.subscribe(topic, new RebalanceListener(rebalanceFlag));
//		consumer.subscribe(topic);
	}
	
	public void connectPartition(List<Integer> partitionList){
		List<TopicPartition> pList = new ArrayList<>();
		for(int partitionId : partitionList){
			TopicPartition tPartition = new TopicPartition(topic.get(0), partitionId);
			pList.add(tPartition);
		}
		consumer.assign(pList);
		consumer = new org.apache.kafka.clients.consumer.KafkaConsumer<String, byte[]>(getConsumerProps());
		consumer.subscribe(topic, new RebalanceListener(rebalanceFlag));
	}
	
	public Iterator<ConsumerRecord<String, byte[]>> poll(){
		ConsumerRecords<String, byte[]> records = consumer.poll(batchDurationMillis);
		if (rebalanceFlag.get()) {
            rebalanceFlag.set(false);
        }
		if(records.isEmpty()){
			try {
				Thread.sleep(500);
			} catch (InterruptedException e) {
				LOG.info("---------kafkaconsumer_Exception----------", e);
			}
		}
//		consumer.commitAsync();
		return records.iterator();
	}

	private Properties getConsumerProps() {
		Properties props = new Properties();
		props.put("bootstrap.servers", brokerList);
		props.put("group.id", groupId);
		props.put("enable.auto.commit", "true");// 自动commit
		props.put("auto.commit.interval.ms", "1500");// 定时commit的周期
		props.put("session.timeout.ms", "30000");// consumer心跳超时时间
		props.put("auto.offset.reset", "latest");
//		props.put("max.poll.interval.ms", "350000");300s
//		props.put("max.poll.records", "200");
		props.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");// 配置序列化类
		props.put("value.deserializer", "org.apache.kafka.common.serialization.ByteArrayDeserializer");
		return props;
	}
	
	public void close(){
//		consumer.commitAsync();
		consumer.wakeup();
	    consumer.close();
	}

	class RebalanceListener implements ConsumerRebalanceListener {

		private AtomicBoolean rebalanceFlag;

		public RebalanceListener(AtomicBoolean rebalanceFlag) {
			this.rebalanceFlag = rebalanceFlag;
		}

		public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
			for (TopicPartition partition : partitions) {
				LOG.info("topic {} - partition {} revoked.", partition.topic(), partition.partition());
				rebalanceFlag.set(true);
			}
		}

		public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
			for (TopicPartition partition : partitions) {
				LOG.info("topic {} - partition {} assigned.", partition.topic(), partition.partition());
			}
		}
	}

	public static void main(String[] args) {
		KafkaConsumer consumer = new KafkaConsumer(args[0],args[1]);
		consumer.setBatchDurationMillis(100L);
		List<String> topics = StringUtil.split(args[2], ",");
		consumer.setTopic(topics);
		consumer.connect();
		int count = 0;
		while(true) {
			Iterator<ConsumerRecord<String, byte[]>> it = consumer.poll();
			while(it.hasNext()) {
				try{
					String msg = new String(it.next().value());
					System.out.println(msg);
					Thread.sleep(5000);
					count++;
					if (count > 5) {
						break;
					}
				} catch (Exception e) {
					e.printStackTrace();
				}
			}
			if (count > 5) {
				break;
			}
		}
		
		consumer.close();
	}
}
